import logging
import os

import matplotlib.pyplot as plt

# import sentencepiece as spm
from transformers import AutoTokenizer


def set_seed(seed):
    import random

    import numpy as np
    import torch

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)


def valid_line(line: str) -> bool:
    invalid_set = {"()", ",,", "..", "''", '""', "“”", "，，", "。。", "（）"}
    for invalid_token in invalid_set:
        if invalid_token in line:
            return False
    return True


def chinese_tokenizer_load():
    # sp_chn = spm.SentencePieceProcessor()
    # sp_chn.Load('{}.model'.format("./tokenizer/chn"))
    # return sp_chn
    model_checkpoint = "Helsinki-NLP/opus-mt-zh-en"
    tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
    return tokenizer


def english_tokenizer_load():
    # sp_eng = spm.SentencePieceProcessor()
    # sp_eng.Load('{}.model'.format("./tokenizer/eng"))
    # return sp_eng
    model_checkpoint = "Helsinki-NLP/opus-mt-en-zh"
    tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
    # tokenizer.tokenize()
    return tokenizer


def set_logger(log_path):
    """Set the logger to log info in terminal and file `log_path`.
    In general, it is useful to have a logger so that every output to the terminal is saved
    in a permanent file. Here we save it to `model_dir/train.log`.
    Example:
    ```
    logging.info("Starting training...")
    ```
    Args:
        log_path: (string) where to log
    """
    if os.path.exists(log_path) is True:
        os.remove(log_path)
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)

    if not logger.handlers:
        # Logging to a file
        file_handler = logging.FileHandler(log_path)
        file_handler.setFormatter(logging.Formatter("%(asctime)s:%(levelname)s: %(message)s"))
        logger.addHandler(file_handler)

        # Logging to console
        stream_handler = logging.StreamHandler()
        stream_handler.setFormatter(logging.Formatter("%(message)s"))
        logger.addHandler(stream_handler)


def draw_training(train_loss_list, valid_loss_list, pic_name):
    num_epochs = len(train_loss_list)
    fig, ax = plt.subplots()
    ax.set_title("Training and Validation Loss")
    ax.set_xlabel("Epochs")
    ax.set_ylabel("Loss")
    ax.plot(range(num_epochs), train_loss_list, label="train loss")
    ax.plot(range(num_epochs), valid_loss_list, label="valid loss")
    ax.legend()
    fig.savefig(pic_name)
    plt.close(fig)
